nanopyx.methods.drift_alignment.estimator_3d

 1import numpy as np
 2from math import sqrt
 3from scipy.interpolate import interp1d
 4
 5from .estimator_table import DriftEstimatorTable
 6from .corrector import DriftCorrector
 7from .estimator import DriftEstimator
 8from ...core.analysis.estimate_shift import GetMaxOptimizer
 9from ...core.utils.timeit import timeit
10from ...core.analysis.ccm import calculate_ccm
11
12
13class Estimator3D(object):
14    """
15    Main class implementing 3d drift correction.
16    Requires an image array with shape (t, z, y, x)
17    """
18
19    def __init__(self):
20        self.image_array = None
21        self.xy_estimator = None
22        self.z_estimator = None
23
24    def correct_xy_drift(self, projection_mode="Mean", **kwargs):
25
26        self.xy_estimator = DriftEstimator()
27
28        if projection_mode == "Mean":
29            projection = np.mean(self.image_array, axis=1)
30        elif projection_mode == "Max":
31            projection = np.max(self.image_array, axis=1)
32        else:
33            print("Not a valid projection mode")
34            return None
35
36        self.xy_estimator.estimate(projection, apply=False, **kwargs)
37
38        corrector = DriftCorrector()
39        corrector.estimator_table = self.xy_estimator.estimator_table
40        for i in range(self.image_array.shape[1]):
41            self.image_array[:, i, :, :] = corrector.apply_correction(self.image_array[:, i, :, :])
42
43    def correct_z_drift(self, axis_mode="top", projection_mode="Mean", **kwargs):
44        
45        if axis_mode == "top":
46            axis_idx = 2
47        elif axis_mode == "left":
48            axis_idx = 3
49        else:
50            print("Not a valid axis mode")
51            return None
52        
53        self.z_estimator = DriftEstimator()
54        if projection_mode == "Mean":
55            projection = np.mean(self.image_array, axis=axis_idx)
56        elif projection_mode == "Max":
57            projection = np.max(self.image_array, axis=axis_idx)
58        else:
59            print("Not a valid projection mode")
60            return None
61
62        self.z_estimator.estimate(projection, apply=False, **kwargs)
63        
64        corrector = DriftCorrector()
65        print(self.image_array.shape, projection.shape)
66        corrector.estimator_table = self.z_estimator.estimator_table
67        if axis_mode == "top":
68            corrector.estimator_table.drift_table[:, 1] = 0
69            for i in range(self.image_array.shape[axis_idx]):
70                self.image_array[:, :, i, :] = corrector.apply_correction(self.image_array[:, :, i, :])
71        elif axis_mode == "left":
72            corrector.estimator_table.drift_table[:, 2] = 0
73            for i in range(self.image_array.shape[axis_idx]):
74                self.image_array[:, :, :, i] = corrector.apply_correction(self.image_array[:, :, :, i])
75                
76    def correct_3d_drift(self, image_array, axis_mode="top", projection_mode="Mean", **kwargs):
77        
78        self.image_array = image_array
79        self.correct_xy_drift(projection_mode=projection_mode, **kwargs)
80        self.correct_z_drift(axis_mode=axis_mode, projection_mode=projection_mode, **kwargs)
81        return self.image_array
82    
class Estimator3D:
14class Estimator3D(object):
15    """
16    Main class implementing 3d drift correction.
17    Requires an image array with shape (t, z, y, x)
18    """
19
20    def __init__(self):
21        self.image_array = None
22        self.xy_estimator = None
23        self.z_estimator = None
24
25    def correct_xy_drift(self, projection_mode="Mean", **kwargs):
26
27        self.xy_estimator = DriftEstimator()
28
29        if projection_mode == "Mean":
30            projection = np.mean(self.image_array, axis=1)
31        elif projection_mode == "Max":
32            projection = np.max(self.image_array, axis=1)
33        else:
34            print("Not a valid projection mode")
35            return None
36
37        self.xy_estimator.estimate(projection, apply=False, **kwargs)
38
39        corrector = DriftCorrector()
40        corrector.estimator_table = self.xy_estimator.estimator_table
41        for i in range(self.image_array.shape[1]):
42            self.image_array[:, i, :, :] = corrector.apply_correction(self.image_array[:, i, :, :])
43
44    def correct_z_drift(self, axis_mode="top", projection_mode="Mean", **kwargs):
45        
46        if axis_mode == "top":
47            axis_idx = 2
48        elif axis_mode == "left":
49            axis_idx = 3
50        else:
51            print("Not a valid axis mode")
52            return None
53        
54        self.z_estimator = DriftEstimator()
55        if projection_mode == "Mean":
56            projection = np.mean(self.image_array, axis=axis_idx)
57        elif projection_mode == "Max":
58            projection = np.max(self.image_array, axis=axis_idx)
59        else:
60            print("Not a valid projection mode")
61            return None
62
63        self.z_estimator.estimate(projection, apply=False, **kwargs)
64        
65        corrector = DriftCorrector()
66        print(self.image_array.shape, projection.shape)
67        corrector.estimator_table = self.z_estimator.estimator_table
68        if axis_mode == "top":
69            corrector.estimator_table.drift_table[:, 1] = 0
70            for i in range(self.image_array.shape[axis_idx]):
71                self.image_array[:, :, i, :] = corrector.apply_correction(self.image_array[:, :, i, :])
72        elif axis_mode == "left":
73            corrector.estimator_table.drift_table[:, 2] = 0
74            for i in range(self.image_array.shape[axis_idx]):
75                self.image_array[:, :, :, i] = corrector.apply_correction(self.image_array[:, :, :, i])
76                
77    def correct_3d_drift(self, image_array, axis_mode="top", projection_mode="Mean", **kwargs):
78        
79        self.image_array = image_array
80        self.correct_xy_drift(projection_mode=projection_mode, **kwargs)
81        self.correct_z_drift(axis_mode=axis_mode, projection_mode=projection_mode, **kwargs)
82        return self.image_array

Main class implementing 3d drift correction. Requires an image array with shape (t, z, y, x)

image_array
xy_estimator
z_estimator
def correct_xy_drift(self, projection_mode='Mean', **kwargs):
25    def correct_xy_drift(self, projection_mode="Mean", **kwargs):
26
27        self.xy_estimator = DriftEstimator()
28
29        if projection_mode == "Mean":
30            projection = np.mean(self.image_array, axis=1)
31        elif projection_mode == "Max":
32            projection = np.max(self.image_array, axis=1)
33        else:
34            print("Not a valid projection mode")
35            return None
36
37        self.xy_estimator.estimate(projection, apply=False, **kwargs)
38
39        corrector = DriftCorrector()
40        corrector.estimator_table = self.xy_estimator.estimator_table
41        for i in range(self.image_array.shape[1]):
42            self.image_array[:, i, :, :] = corrector.apply_correction(self.image_array[:, i, :, :])
def correct_z_drift(self, axis_mode='top', projection_mode='Mean', **kwargs):
44    def correct_z_drift(self, axis_mode="top", projection_mode="Mean", **kwargs):
45        
46        if axis_mode == "top":
47            axis_idx = 2
48        elif axis_mode == "left":
49            axis_idx = 3
50        else:
51            print("Not a valid axis mode")
52            return None
53        
54        self.z_estimator = DriftEstimator()
55        if projection_mode == "Mean":
56            projection = np.mean(self.image_array, axis=axis_idx)
57        elif projection_mode == "Max":
58            projection = np.max(self.image_array, axis=axis_idx)
59        else:
60            print("Not a valid projection mode")
61            return None
62
63        self.z_estimator.estimate(projection, apply=False, **kwargs)
64        
65        corrector = DriftCorrector()
66        print(self.image_array.shape, projection.shape)
67        corrector.estimator_table = self.z_estimator.estimator_table
68        if axis_mode == "top":
69            corrector.estimator_table.drift_table[:, 1] = 0
70            for i in range(self.image_array.shape[axis_idx]):
71                self.image_array[:, :, i, :] = corrector.apply_correction(self.image_array[:, :, i, :])
72        elif axis_mode == "left":
73            corrector.estimator_table.drift_table[:, 2] = 0
74            for i in range(self.image_array.shape[axis_idx]):
75                self.image_array[:, :, :, i] = corrector.apply_correction(self.image_array[:, :, :, i])
def correct_3d_drift(self, image_array, axis_mode='top', projection_mode='Mean', **kwargs):
77    def correct_3d_drift(self, image_array, axis_mode="top", projection_mode="Mean", **kwargs):
78        
79        self.image_array = image_array
80        self.correct_xy_drift(projection_mode=projection_mode, **kwargs)
81        self.correct_z_drift(axis_mode=axis_mode, projection_mode=projection_mode, **kwargs)
82        return self.image_array